
#ifndef AFTERMATH_NUMERICS_VSSE_MATH_HPP_INCLUDED
#define AFTERMATH_NUMERICS_VSSE_MATH_HPP_INCLUDED

#include "inline.hpp"

#include <immintrin.h>
#include <mmintrin.h>
#include <smmintrin.h>

namespace aftermath
{
    namespace numerics
    {
        template <int mask>
        OP_INLINE __m256d _mm256_dp_pd(const __m256d& left, const __m256d& right)
        {
            if ((mask & 0xF0) == 0x00) return _mm256_setzero_pd();

            // broadcast if m0 == m2 and m1 == m3
            if (((mask & 0x1) == ((mask & 0x4) >> 2)) && ((mask & 0x2) == ((mask & 0x8) >> 2)))
            {
                __m128d a, b, c, d;
                a = _mm256_extractf128_pd(left, 0);
                b = _mm256_extractf128_pd(right, 0);
                // 0-0-m5-m4-0-0-m1-m0: (and 00110011)
                c = _mm_dp_pd(a, b, (mask & 0x33));
            
                a = _mm256_extractf128_pd(left, 1);
                b = _mm256_extractf128_pd(right, 1);
                // 0-0-m7-m6-0-0-m1-m0: (and 11000000, >> 2) + (and 00000011)
                d = _mm_dp_pd(a, b, ((mask & 0xC0) >> 2) + (mask & 0x3));

                a = _mm_add_pd(c, d);
                return _mm256_broadcast_pd(&a);
            }
            // simple insert when (m1, m0) = 0 or (m3, m2) = 0
            // then to choose the nonzero one just take (mask & 00000011) + ((mask & 00001100) >> 2)
            else if (((mask & 0x3) == 0) || ((mask & 0xC) == 0))
            {
                __m128d a, b, c, d;
                a = _mm256_extractf128_pd(left, 0);
                b = _mm256_extractf128_pd(right, 0);
                // 0-0-m5-m4-0-0-?-?: (and 00110000) + (??)
                c = _mm_dp_pd(a, b, (mask & 0x30) + (mask & 0x3) + ((mask & 0xC) >> 2));
            
                a = _mm256_extractf128_pd(left, 1);
                b = _mm256_extractf128_pd(right, 1);
                // 0-0-m7-m6-0-0-?-?: (and 11000000, >> 2) + (??)
                d = _mm_dp_pd(a, b, ((mask & 0xC0) >> 2) + (mask & 0x3) + ((mask & 0xC) >> 2));

                a = _mm_add_pd(c, d);
                // if (m1, m0) = 0 then offset = 1, if (m3, m2) = 0 then offset is 0
                return _mm256_insertf128_pd(_mm256_setzero_pd(), a, ((mask & 0x3) == 0) ? 1 : 0);
            }
            else 
            {
                __m256d a, b;
                a = _mm256_mul_pd(left, right);
                // zero components that do not participate in dot product
                if (mask & 0xF0 != 0xF0) a = _mm256_blend_pd(_mm256_setzero_pd(), a, (mask >> 4));

                b = _mm256_permute2f128_pd(a, a, 0x1); // b = a1-a0-a3-a2
                a = _mm256_add_pd(a, b); // a = [a1 + a3, a0 + a2, a3 + a1, a2 + a0]
                if (mask & 0xF == 0xF) return _mm256_hadd_pd(a, a);
                else return _mm256_blend_pd(_mm256_setzero_pd(), _mm256_hadd_pd(a, a), (mask & 0xF));
            }
        }

        OP_INLINE __m256d operator +(const __m256d& left, const __m256d& right) { return _mm256_add_pd(left, right); }
        OP_INLINE __m256d operator -(const __m256d& left, const __m256d& right) { return _mm256_sub_pd(left, right); }
        OP_INLINE __m256d operator *(const __m256d& left, const __m256d& right) { return _mm256_mul_pd(left, right); }
        OP_INLINE __m256d operator /(const __m256d& left, const __m256d& right) { return _mm256_div_pd(left, right); }
    
        OP_INLINE __m256d& operator +=(__m256d& left, const __m256d& right) { return left = _mm256_add_pd(left, right); }
        OP_INLINE __m256d& operator -=(__m256d& left, const __m256d& right) { return left = _mm256_sub_pd(left, right); }
        OP_INLINE __m256d& operator *=(__m256d& left, const __m256d& right) { return left = _mm256_mul_pd(left, right); }
        OP_INLINE __m256d& operator /=(__m256d& left, const __m256d& right) { return left = _mm256_div_pd(left, right); }

        template <int nRows, int nColumns>
        OP_INLINE void quad_row_transpose(
            const __m256d& row0, const __m256d& row1, const __m256d& row2, const __m256d& row3,
            __m256d& col0, __m256d& col1, __m256d& col2, __m256d& col3,
            __m256d& tmp0, __m256d& tmp1, __m256d& tmp2, __m256d& tmp3, __m256d& zero) 
        {
            if (nRows == 0 || nColumns == 0) return;
            if (nRows == 1 && nColumns == 1) 
            {
                col0 = row0;
                return;
            }

            // most significant <------ least significant
            // row0: a3 a2 a1 a0
            // row1: b3 b2 b1 b0
            // row2: c3 c2 c1 c0
            // row3: d3 d2 d1 d0

            tmp0 = _mm256_permute2f128_pd(row0, nRows > 2 ? row2 : zero, 0x20); // tmp0: c1 c0 a1 a0
            if (nColumns > 2) tmp2 = _mm256_permute2f128_pd(row0, nRows > 2 ? row2 : zero, 0x31); // tmp2: c3 c2 a3 a2
            if (nRows > 1)
            {
                tmp1 = _mm256_permute2f128_pd(row1, nRows > 3 ? row3 : zero, 0x20); // tmp1: d1 d0 b1 b0
                if (nColumns > 2) tmp3 = _mm256_permute2f128_pd(row1, nRows > 3 ? row3 : zero, 0x31); // tmp3: d3 d2 b3 b2
            }
        
            if (nColumns > 3) col3 = _mm256_unpackhi_pd(tmp2, nRows > 1 ? tmp3 : zero); // col3: d3 c3 b3 a3
            if (nColumns > 2) col2 = _mm256_unpacklo_pd(tmp2, nRows > 1 ? tmp3 : zero); // col2: d2 c2 b2 a2
            if (nColumns > 1) col1 = _mm256_unpackhi_pd(tmp0, nRows > 1 ? tmp1 : zero); // col1: d1 c1 b1 a1
            col0 = _mm256_unpacklo_pd(tmp0, nRows > 1 ? tmp1 : zero); // col0: d0 c0 b0 a0
        }

        template <int nRows>
        OP_INLINE void quad_row_transpose_self(
            __m256d& row0, __m256d& row1, __m256d& row2, __m256d& row3,
            __m256d& tmp0, __m256d& tmp2, __m256d& zero) 
        {
            if (nRows == 0) return;
            if (nRows == 1) return;

            // most significant <------ least significant
            // row0: a3 a2 a1 a0
            // row1: b3 b2 b1 b0
            // row2: c3 c2 c1 c0
            // row3: d3 d2 d1 d0
        
            tmp0 = _mm256_permute2f128_pd(row0, nRows > 2 ? row2 : zero, 0x20); // tmp0: c1 c0 a1 a0
            if (nRows > 2) tmp2 = _mm256_permute2f128_pd(row0, row2, 0x31); // tmp2: c3 c2 a3 a2
            // We don't need row0 and row2 any more; might as well use them
            if (nRows > 1)
            {
                row0 = _mm256_permute2f128_pd(row1, nRows > 3 ? row3 : zero, 0x20); // row0: d1 d0 b1 b0
                // row0 (not row2!) now stores what used to be tmp1
                if (nRows > 2) row2 = _mm256_permute2f128_pd(row1, nRows > 3 ? row3 : zero, 0x31); // row2: d3 d2 b3 b2
                // row2 now stores what used to be tmp3, we will get rid of it soon
                // Now we don't need original row1 and row3, but we can't overwrite row0 or row2 just yet; so start with the former
            }
        
            if (nRows > 3) row3 = _mm256_unpackhi_pd(tmp2, nRows > 1 ? row2 : zero); // row3: d3 c3 b3 a3
            // now we don't need row2 (former tmp3): overwrite it
            if (nRows > 2) row2 = _mm256_unpacklo_pd(tmp2, nRows > 1 ? row2 : zero); // row2: d2 c2 b2 a2
            if (nRows > 1) row1 = _mm256_unpackhi_pd(tmp0, nRows > 1 ? row0 : zero); // row1: d1 c1 b1 a1
            // finally overwrite row0
            row0 = _mm256_unpacklo_pd(tmp0, nRows > 1 ? row0 : zero); // row0: d0 c0 b0 a0
        }

        template <int nRows>
        OP_INLINE void quad_coupled_row_transpose(
            __m256d& urow0, __m256d& urow1, __m256d& urow2, __m256d& urow3,
            __m256d& lrow0, __m256d& lrow1, __m256d& lrow2, __m256d& lrow3,
            __m256d& utmp0, __m256d& utmp1, __m256d& utmp2, __m256d& utmp3,
            __m256d& ltmp0, __m256d& ltmp1, __m256d& ltmp2, __m256d& ltmp3, __m256d& zero) 
        {
            if (nRows == 0) return;
            // For upper: nRows = 4, nColumns = nRows
            // For lower: nRows = nRows, nColumns = 4

            utmp0 = _mm256_permute2f128_pd(urow0, urow2, 0x20);
            if (nRows > 2) utmp2 = _mm256_permute2f128_pd(urow0, urow2, 0x31);
            utmp1 = _mm256_permute2f128_pd(urow1, urow3, 0x20);
            if (nRows > 2) utmp3 = _mm256_permute2f128_pd(urow1, urow3, 0x31);

            quad_row_transpose<nRows, 4>(
                lrow0, lrow1, lrow2, lrow3,
                urow0, urow1, urow2, urow3,
                ltmp0, ltmp1, ltmp2, ltmp3, zero
                );
            //ltmp0 = _mm256_permute2f128_pd(lrow0, nRows > 2 ? lrow2 : zero, 0x20);
            //ltmp2 = _mm256_permute2f128_pd(lrow0, nRows > 2 ? lrow2 : zero, 0x31);
            //if (nRows > 1)
            //{
            //    ltmp1 = _mm256_permute2f128_pd(lrow1, nRows > 3 ? lrow3 : zero, 0x20);
            //    ltmp3 = _mm256_permute2f128_pd(lrow1, nRows > 3 ? lrow3 : zero, 0x31);
            //}

            //urow0 = _mm256_unpacklo_pd(ltmp0, nRows > 1 ? ltmp1 : zero);
            //urow1 = _mm256_unpackhi_pd(ltmp0, nRows > 1 ? ltmp1 : zero);
            //urow2 = _mm256_unpacklo_pd(ltmp2, nRows > 1 ? ltmp3 : zero);
            //urow3 = _mm256_unpackhi_pd(ltmp2, nRows > 1 ? ltmp3 : zero);
        
            lrow0 = _mm256_unpacklo_pd(utmp0, utmp1);
            if (nRows > 1) lrow1 = _mm256_unpackhi_pd(utmp0, utmp1);
            if (nRows > 2) lrow2 = _mm256_unpacklo_pd(utmp2, utmp3);
            if (nRows > 3) lrow3 = _mm256_unpackhi_pd(utmp2, utmp3);
        }

        template <int nRows, int nColumns, bool doIncrement>
        OP_INLINE void quad_multiply(
            const __m256d& row0, const __m256d& row1, const __m256d& row2, const __m256d& row3,
            const __m256d& col0, const __m256d& col1, const __m256d& col2, const __m256d& col3,
            __m256d& res0, __m256d& res1, __m256d& res2, __m256d& res3,
            __m256d& tmp0, __m256d& tmp1, __m256d& tmp2, __m256d& tmp3, __m256d& zero) 
        {
            if (nRows == 0 || nColumns == 0) return;

            // We don't need tmp* then!
            if (nColumns == 1)
            {
                if (doIncrement)
                {
                    res0 += _mm256_dp_pd<0xF1>(row0, col0);
                    if (nRows > 1) res1 += _mm256_dp_pd<0xF1>(row1, col0);
                    if (nRows > 2) res2 += _mm256_dp_pd<0xF1>(row2, col0);
                    if (nRows > 3) res3 += _mm256_dp_pd<0xF1>(row3, col0);
                }
                else
                {
                    res0 = _mm256_dp_pd<0xF1>(row0, col0);
                    if (nRows > 1) res1 = _mm256_dp_pd<0xF1>(row1, col0);
                    if (nRows > 2) res2 = _mm256_dp_pd<0xF1>(row2, col0);
                    if (nRows > 3) res3 = _mm256_dp_pd<0xF1>(row3, col0);
                }
            }
            else
            {
                tmp0 = _mm256_dp_pd<0xF1>(row0, col0);
                tmp1 = _mm256_dp_pd<0xF2>(row0, col1);
                if (nColumns > 2) tmp2 = _mm256_dp_pd<0xF4>(row0, col2);
                if (nColumns > 3) tmp3 = _mm256_dp_pd<0xF8>(row0, col3);

                if (doIncrement)
                {
                    res0 += tmp0;
                    res0 += tmp1;
                }
                else res0 = tmp0 + tmp1;
                if (nColumns > 2) res0 += tmp2;
                if (nColumns > 3) res0 += tmp3;

                if (nRows > 1)
                {
                    tmp0 = _mm256_dp_pd<0xF1>(row1, col0);
                    tmp1 = _mm256_dp_pd<0xF2>(row1, col1);
                    if (nColumns > 2) tmp2 = _mm256_dp_pd<0xF4>(row1, col2);
                    if (nColumns > 3) tmp3 = _mm256_dp_pd<0xF8>(row1, col3);
            
                    if (doIncrement)
                    {
                        res1 += tmp0;
                        res1 += tmp1;
                    }
                    else res1 = tmp0 + tmp1;
                    if (nColumns > 2) res1 += tmp2;
                    if (nColumns > 3) res1 += tmp3;
                }
                if (nRows > 2)
                {
                    tmp0 = _mm256_dp_pd<0xF1>(row2, col0);
                    tmp1 = _mm256_dp_pd<0xF2>(row2, col1);
                    if (nColumns > 2) tmp2 = _mm256_dp_pd<0xF4>(row2, col2);
                    if (nColumns > 3) tmp3 = _mm256_dp_pd<0xF8>(row2, col3);
            
                    if (doIncrement)
                    {
                        res2 += tmp0;
                        res2 += tmp1;
                    }
                    else res2 = tmp0 + tmp1;
                    if (nColumns > 2) res2 += tmp2;
                    if (nColumns > 3) res2 += tmp3;
                }
                if (nRows > 3)
                {
                    tmp0 = _mm256_dp_pd<0xF1>(row3, col0);
                    tmp1 = _mm256_dp_pd<0xF2>(row3, col1);
                    if (nColumns > 2) tmp2 = _mm256_dp_pd<0xF4>(row3, col2);
                    if (nColumns > 3) tmp3 = _mm256_dp_pd<0xF8>(row3, col3);
            
                    if (doIncrement)
                    {
                        res3 += tmp0;
                        res3 += tmp1;
                    }
                    else res3 = tmp0 + tmp1;
                    if (nColumns > 2) res3 += tmp2;
                    if (nColumns > 3) res3 += tmp3;
                }
            }
        }
    }
}

#endif // AFTERMATH_NUMERICS_VSSE_MATH_HPP_INCLUDED
